import os
import argparse
import time

from extract import Document
from utils import generator_topic_string

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="提取网页中的题目 用于复习")
    # parser.add_argument("--type", type=str, help="读取单文件或读取文件夹")
    parser.add_argument("--path", type=str, help="读取的文件路径")
    parser.add_argument("--to", type=str, default="txt", help="保存为什么格式")
    # parser.add_argument("--split", type=int, default=0, help="是否分开保存")
    args = parser.parse_args()
    path = args.path
    to = args.to
    # split = args.split
    if path is None:
        path = os.getcwd()
    result = []
    if os.path.isfile(path):
        with open(path, "r", encoding="utf-8") as f:
            html = f.read()
        document = Document(html)
        lst = document.extract()
        topics = generator_topic_string(lst)
        result.extend(topics)
    if os.path.isdir(path):
        for root, dirs, files in os.walk(path):
            for file in files:
                print(file)
                if file.endswith(".html"):
                    with open(os.path.join(root, file), "r", encoding="utf-8") as f:
                        html = f.read()
                    document = Document(html)
                    lst = document.extract()
                    topics = generator_topic_string(lst)
                    result.extend(topics)
    if to == "txt":
        with open("result.txt", "w", encoding="utf-8") as f:
            f.write("\n".join(result))
        print("提取完毕")
    time.sleep(5)
